% Chakram is a subproject of Hexena - Haskell Executable Analyser
% License:			GPL-3
% Maintainer:		Montez Fitzpatrick
% Stability:		Experimental
% Portability:		GHC
% Module:			Hexena.Codec.Parser.MzHeader

Check [http://msdn.microsoft.com/en-us/magazine/cc301805.aspx]
for additional information related to Win32 PE file format

                     ** IMAGE_DOS_HEADER **

--                                     1  1  1  1  1  1
--       0  1  2  3  4  5  6  7  8  9  0  1  2  3  4  5
--     +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
--     |                      magic                    |
--     +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
--     |                      cblp                     |
--     +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
--     |                      cp                       |
--     +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
--     |                      crlc                     |
--     +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
--     |                      cparhdr                  |
--     +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
--     |                      minalloc                 |
--     +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
--     |                      maxalloc                 |
--     +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
--     |                      ss                       |
--     +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
--     |                      sp                       |
--     +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
--     |                      csum                     |
--     +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
--     |                      ip                       |
--     +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
--     |                      cs                       |
--     +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
--     |                      ifarlc                   |
--     +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
--     |                      ovno                     |
--     +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
--     |                      res [4]                  |
--     |                                               |
--     |                                               |
--     |                                               |
--     +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
--     |                      oemid                    |
--     +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
--     |                      oeminfo                  |
--     +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
--     |                      res2 [10]                |
--     |                                               |
--     |                 ** omitted **                 |
--     |                                               |
--     |                                               |
--     +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
--     |                      lfanew                   |
--     |                                               |
--     +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+

\begin{code}
module Hexena.Codec.Parser.MzHeader (parseMzHeader) where

import Data.Binary.Get
import Data.Word

import Hexena.Data.MzHeader (MzHeader(..), Res1(..), Res2(..))

_SIZE_OF_IMAGE_DOS_HEADER :: Int
_SIZE_OF_IMAGE_DOS_HEADER = 64 -- 64 Bytes

-- | Basic header population, probably would be better to use the setter
--   funtions from the lens, that could be a future experimentation
--   and refactorization.
--
--   Usage is with the runGet function. 
--
--   Descriptions derived from WINNT.H
parseMzHeader :: Get MzHeader
parseMzHeader = do
  w_magic <- getWord16le        -- Magic Number
  w_cblp <- getWord16le            -- Bytes on last page of file    
  w_cp <- getWord16le            -- Pages in file
  w_crlc <- getWord16le            -- Relocations
  w_cparhdr <- getWord16le        -- Size of header in paragraphs
  w_minalloc <- getWord16le        -- Minimum extra paragraphs needed
  w_maxalloc <- getWord16le        -- Maximum extra paragraphs needed
  w_ss <- getWord16le            -- Initial (relative) SS value
  w_sp <- getWord16le            -- Initial SP value
  w_csum <- getWord16le            -- Checksum
  w_ip <- getWord16le            -- Initial IP value
  w_cs <- getWord16le            -- Initial (relative) CS value
  w_lrarlc <- getWord16le        -- File address of relocation table
  w_ovno <- getWord16le            -- Overlay number
  w_res <- parseRes1            -- Reserved words
  w_oemid <- getWord16le        -- OEM identifier (for e_oeminfo)
  w_oeminfo <- getWord16le        -- OEM information; e_oemid specific
  w_res2 <- parseRes2            -- Reserved words
  w_lfanew <- getWord32le        -- File address of new exe header
 
  return $ MzHeader w_magic w_cblp w_cp w_crlc w_cparhdr w_minalloc w_maxalloc w_ss w_sp w_csum w_ip w_cs w_lrarlc w_ovno w_res w_oemid w_oeminfo w_res2 w_lfanew
  
-- | Reserved Words Header
parseRes1 :: Get Res1
parseRes1 = do
  r0 <- getWord16le
  r1 <- getWord16le
  r2 <- getWord16le
  r3 <- getWord16le
  return $ Res1 r0 r1 r2 r3  

-- | Reserved Words 2 Header
parseRes2 :: Get Res2
parseRes2 = do
  r0 <- getWord16le
  r1 <- getWord16le
  r2 <- getWord16le
  r3 <- getWord16le
  r4 <- getWord16le
  r5 <- getWord16le
  r6 <- getWord16le
  r7 <- getWord16le
  r8 <- getWord16le
  r9 <- getWord16le
  return $ Res2 r0 r1 r2 r3 r4 r5 r6 r7 r8 r9

-- | Default MZ Header as a Word8 list.
commonDefaultHeaderWord8 :: [Word8]
commonDefaultHeaderWord8 = [ 0x4d, 0x5a, 0x90, 0x00, 0x03, 0x00, 0x00, 0x00, 
                             0x04, 0x00, 0x00, 0x00, 0xff, 0xff, 0x00, 0x00, 
                             0xb8, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 
                             0x40, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 
                             0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 
                             0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 
                             0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 
                             0x00, 0x00, 0x00, 0x00, 0xFF, 0xFF, 0xFF, 0xFF]


-- | Utility function to take a MzHeader and compare it to the common default header and see if
--   there is a match.  This function will overwrite using the same lfanew.  Uses 
--   binary-stict
--
--   Will return a boolean.
isCommonMzHeader hdr = 
    if compare hdr (overwriteLfaNew hdr (runGet parseMzHeader $ L.pack commonDefaultHeaderWord8)) == EQ
        then True
        else False


-- | Binary uses an accumlator for counting what has been consumed.  This just
--   matches '0' since I don't like magic numbers where they aren't strictly
--   obvious.
startOfFileOffset :: Int
startOfFileOffset = 0

-- | Internal function that commit the overwrite from the working executable
--   to the common default header
overwriteLfaNew :: MzHeader -> MzHeader -> MzHeader
overwriteLfaNew fh com = set lfanew (get lfanew fh) com

-- | Internal helper function to unwrap an Right either will send an error on Left
--  
--  TODO: Have this function on Left return ByteString.empty instead of generating an error.
fromRight :: Either a b  -> b 
fromRight (Right b) = b
fromRight _ = error "Did not parse MZ Header correctly." 

-- | Calculate Length to PE Header function calculates how large the jump is from 
--   the current point in the stream to the start of the PE Header.  This is useful
--   when MZHeader has been completely consumed so the DOS Stub space length can
--   be calculated.
--
calcLenToPEH :: Int -> MzHeader -> Int
calcLenToPEH consumed mzhdr = lenToStubEnd - consumed
  where
	lenToStubEnd = fromIntegral $ get lfanew mzhdr
\end{code}
